from tornado.web import Application, RequestHandler
from tornado.ioloop import IOLoop
import json
import joblib
import pandas as pd
from test import load_dataset, getLabelEncode, full_pipeline


class HelloHandler(RequestHandler):
    def get(self):
        self.write({'message': 'Hello from MLServer'})


# ML based classification
class MLHandler(RequestHandler):
    def post(self):
        # print(json.loads(self.request.body))
        # print(self.request.body)
        dic = json.loads(self.request.body)
        X = pd.DataFrame(dic)
        X_set = X.drop("class", axis=1)
        test_set_prepared = full_pipeline(X_set, num_attribs, cat_attribs, kddcup)
        result = clf.predict(test_set_prepared)
        # print("result:", result[0])
        if result[0]:
            print("attack stream detected")
            print(dic)
            self.write({'message': "positive"})
        else:
            print("normal stream")
            self.write({'message': "negative"})


# Reinforcement Learning
class RLHandler(RequestHandler):
    def post(self):
        print(json.loads(self.request.body))
        # print(self.request.body)
        # TODO: reinforcement learning
        self.write({'message': 'RL'})


def make_app():
    urls = [
        ("/", HelloHandler),
        ("/ML/", MLHandler),
        ("/RL/", RLHandler)
    ]
    return Application(urls, debug=True)


if __name__ == '__main__':
    # ML test
    dic = {'duration': [0], 'protocol_type': 'tcp', 'service': 'http', 'flag': 'SF', 'src_bytes': 181,
           'dst_bytes': 5450, 'land': 0, 'wrong_fragment': 0, 'urgent': 0, 'ho': 0, 'num_failed_logins': 0,
           'logged_in': 1,
           'num_compromised': 0, 'root_shell': 0, 'su_attempted': 0, 'num_root': 0, 'num_file_creations': 0,
           'num_shells': 0, 'num_access_files': 0, 'num_outbound_cmds': 0, 'is_host_login': 0, 'is_guest_login': 0,
           'count': 8, 'srv_count': 8, 'serror_rate': 0.0, 'srv_serror_rate': 0.0, 'rerror_rate': 0.0,
           'srv_rerror_rate': 0.0, 'same_srv_rate': 1.0, 'diff_srv_rate': 0.0, 'srv_diff_host_rate': 0.0,
           'dst_host_count': 9, 'dst_host_srv_count': 9, 'dst_host_same_srv_rate': 1.0, 'dst_host_diff_srv_rate': 0.0,
           'dst_host_same_src_port_rate': 0.11, 'dst_host_srv_diff_host_rate': 0.0, 'dst_host_serror_rate': 0.0,
           'dst_host_srv_serror_rate': 0.0, 'dst_host_rerror_rate': 0.0, 'dst_host_srv_rerror_rate': 0.0,
           'class': 'normal.'}
    X = pd.DataFrame(dic)
    X_set = X.drop("class", axis=1)

    kddcup = load_dataset()
    labels, NORMAL_NUM = getLabelEncode(kddcup, ["class"], kddcup)
    kddcup["class"] = labels
    num_attribs = list(kddcup)  # 42 - 3 - 1 = 38
    num_attribs.remove("class")
    cat_attribs = ["protocol_type", "service", "flag"]
    for ele in cat_attribs:
        num_attribs.remove(ele)

    MODEL_NAME = "train/test3/rdm_forest_clf.pkl"
    clf = joblib.load(MODEL_NAME)
    test_set_prepared = full_pipeline(X_set, num_attribs, cat_attribs, kddcup)
    result = clf.predict(test_set_prepared)
    print("test result:", result[0])

    # Web server 
    app = make_app()
    app.listen(3000)
    print("listening on port 3000")
    IOLoop.instance().start()
